# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# The `Tensor` template `tensor.prototype.pyi` for `tools/gen_tensor_stub.py` to generate the stub file `tensor.pyi`.
# Add docstring, attributes, methods and alias with type annotations for `Tensor` in `tensor.prototype.pyi`
# if not conveniently coding in original place (like c++ source file).

# Import common typings for generated methods
# isort: off
from typing import *  # noqa: F403
from typing_extensions import *  # type: ignore # noqa: F403
from paddle._typing import *  # noqa: F403

# isort: on
from builtins import (  # noqa: F401
    bool as _bool,
    bytes as _bytes,
    complex as _complex,
    float as _float,
    int as _int,
    str as _str,
)
from collections.abc import Iterator
from typing import Any, Literal, overload

import numpy.typing as npt

import paddle
from paddle import (
    ParamAttr,  # noqa: F401
    _typing,
)
from paddle.base.dygraph.tensor_patch_methods import (
    TensorHookRemoveHelper,  # noqa: F401
)
from paddle.tensor.linalg import _POrder  # noqa: F401
from paddle.tensor.stat import _Interpolation  # noqa: F401

# annotation: ${eager_param_base_begin}
class AbstractEagerParamBase:
    # annotation: ${eager_param_base_docstring}

    # annotation: ${eager_param_base_attributes}

    # annotation: ${eager_param_base_methods}
    @property
    def trainable(self) -> _bool: ...
    @trainable.setter
    def trainable(self, trainable: _bool) -> None: ...

    # annotation: ${eager_param_base_alias}

# annotation: ${eager_param_base_end}

# annotation: ${tensor_begin}
class AbstractTensor:
    # annotation: ${tensor_attributes}

    # If method defined below, we should make the method's signature complete,
    # and ignore the signature extracted from `paddle.Tensor`.
    # `gen_tensor.stub.py` will NOT overwrite the signature below.
    # If method has docstring (ignoring the spaces), `gen_tensor.stub.py` also will NOT overwrite it.

    # annotation: ${tensor_methods}
    @overload
    def __init__(self) -> None: ...
    @overload
    def __init__(
        self, dtype, dims, name: _str, type, persistable: _bool
    ) -> None: ...
    @overload
    def __init__(
        self,
        value: npt.NDArray[Any],
        place,
        persistable: _bool,
        zero_copy: _bool,
        name: _str,
        stop_gradient: _bool,
    ) -> None: ...
    @overload
    def __init__(self, value: npt.NDArray[Any]) -> None: ...
    @overload
    def __init__(
        self, value: Tensor, dims, name: _str, process_mesh, placements
    ) -> None: ...
    @overload
    def __init__(self, value: Tensor, place, name: _str) -> None: ...
    @overload
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        """
        ref: paddle/fluid/pybind/eager.cc

        We should have init function with signature:
        1.
        def __init__ ()
        2.
        def __init__ (
            dtype: paddle::framework::proto::VarType::Type,
            dims: vector<int>,
            name: std::string,
            type: paddle::framework::proto::VarType::DenseTensor,
            persistable: bool)
        3. (multi-place)
        (should have at least one parameter, one parameter equals to case 4, zero
        parameter equals to case 1)
        def __init__ (
            value: ndarray,
            place: paddle::platform::Place,
            persistable: bool,
            zero_copy: bool,
            name: std::string,
            stop_gradient: bool)
        4.
        def __init__ (
            value: ndarray)
        5.
        def __init__ (
            tensor: Tensor)
        6. (multi-place)
        (should have at least one parameter, one parameter equals to case 5, zero
        parameter equals to case 1.)
        def __init__ (
            global_tensor: Tensor,
            place: paddle::platform::Place,
            name: std::string,
            process_mesh: phi::distributed::ProcessMesh
            placements: std::vector<Placement>)
        7. (multi-place)
        (should have at least one parameter, one parameter equals to case 5, zero
        parameter equals to case 1.)
        def __init__ (
            local_tensor: Tensor,
            global_dims: vector<int>,
            name: std::string,
            process_mesh: phi::distributed::ProcessMesh
            placements: std::vector<Placement>)
        8. (multi-place) (should have at least one parameter, one parameter similar
        to case 5, zero parameter equals to case 1.)
        def __init__ (
            tensor: FrameworkTensor,
            place: paddle::platform::Place,
            name: std::string)
        """
        ...
    # rich comparison
    def __eq__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore[override]
    def __ge__(self, y: _typing.TensorLike) -> Tensor: ...
    def __gt__(self, y: _typing.TensorLike) -> Tensor: ...
    def __lt__(self, y: _typing.TensorLike) -> Tensor: ...
    def __le__(self, y: _typing.TensorLike) -> Tensor: ...
    def __ne__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore[override]

    # binary arithmetic operations
    def __add__(self, y: _typing.TensorLike) -> Tensor: ...
    def __sub__(self, y: _typing.TensorLike) -> Tensor: ...
    def __mul__(self, y: _typing.TensorLike) -> Tensor: ...
    def __matmul__(self, y: _typing.TensorLike) -> Tensor: ...
    def __truediv__(self, y: _typing.TensorLike) -> Tensor: ...
    def __floordiv__(self, y: _typing.TensorLike) -> Tensor: ...
    def __mod__(self, y: _typing.TensorLike) -> Tensor: ...
    def __pow__(self, y: _typing.TensorLike) -> Tensor: ...
    def __and__(self, y: _typing.TensorLike) -> Tensor: ...
    def __ror__(self, y: _typing.TensorLike) -> Tensor: ...
    def __rxor__(self, y: _typing.TensorLike) -> Tensor: ...
    def __div__(self, y: _typing.TensorLike) -> Tensor: ...
    def __radd__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
    def __rsub__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
    def __rmul__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
    def __rmatmul__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
    def __rtruediv__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
    def __rmod__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
    def __rpow__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
    def __rdiv__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
    def __rfloordiv__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
    def __rand__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore

    # type cast
    def __bool__(self) -> _bool: ...
    def __float__(self) -> _float: ...
    def __int__(self) -> _int: ...
    def __long__(self) -> _float: ...
    def __nonzero__(self) -> _bool: ...
    def __complex__(self) -> _complex: ...

    # emulating container types
    def __getitem__(
        self,
        item: _typing.TensorIndex,
    ) -> Tensor: ...
    def __setitem__(
        self,
        item: _typing.TensorIndex,
        value: Tensor | npt.NDArray[Any] | _complex | _bool,
    ) -> None: ...
    def __len__(self) -> _int: ...

    # emulating numeric types
    def __index__(self) -> _int: ...

    # unary arithmetic operations
    def __invert__(self) -> Tensor: ...
    def __neg__(self) -> Tensor: ...
    def __pos__(self) -> Tensor: ...

    # basic
    def __hash__(self) -> _int: ...
    def clear_gradient(self, set_to_zero: _bool = True) -> None: ...
    def clone(self) -> Tensor: ...
    def cols(self) -> Tensor: ...
    def contiguous(self) -> Tensor: ...
    def copy_(self) -> Tensor: ...
    def crows(self) -> Tensor: ...
    @property
    def data(self) -> Tensor: ...
    @data.setter
    def data(self, value: Tensor) -> None: ...
    def data_ptr(self) -> _int: ...
    def dense_dim(self) -> _int: ...
    def detach(self) -> Tensor: ...
    def detach_(self) -> Tensor: ...
    @property
    def dtype(self) -> paddle.dtype: ...
    def element_size(self) -> _int: ...
    def get_map_tensor(self) -> Tensor: ...
    def get_selected_rows(self) -> None: ...
    def get_strides(self) -> list[_int]: ...
    def get_tensor(self) -> Tensor: ...
    @property
    def grad(self) -> Tensor | None: ...
    @grad.setter
    def grad(self, value: Tensor | None) -> None: ...
    @property
    def grad_(self) -> Tensor | None: ...
    @grad_.setter
    def grad_(self, value: Tensor) -> None: ...
    @property
    def grad_fn(self) -> Any: ...
    def is_contiguous(self) -> _bool: ...
    def is_coalesced(self) -> _bool: ...
    def is_dense(self) -> _bool: ...
    def is_dist(self) -> _bool: ...
    @property
    def is_leaf(self) -> _bool: ...
    def is_same_shape(self, y: Tensor) -> _bool: ...
    def is_selected_rows(self) -> _bool: ...
    def is_sparse(self) -> _bool: ...
    def is_sparse_coo(self) -> _bool: ...
    def is_sparse_csr(self) -> _bool: ...
    @property
    def layout(self) -> _typing.DataLayoutND: ...
    @property
    def name(self) -> _str: ...
    @name.setter
    def name(self, value: _str) -> None: ...
    @property
    def ndim(self) -> _int: ...
    def nnz(self) -> _int: ...
    @property
    def num_shard(self) -> _int: ...
    def numpy(self) -> npt.NDArray[Any]: ...
    @property
    def offset(self) -> _int: ...
    @property
    def persistable(self) -> _bool: ...
    @persistable.setter
    def persistable(self, value: _bool) -> None: ...
    @property
    def place(self) -> paddle.core.Place: ...
    @property
    def placements(self) -> list[paddle.distributed.Placement] | None: ...
    @property
    def process_mesh(self) -> paddle.distributed.ProcessMesh | None: ...
    def rows(self) -> list[_int]: ...
    def set_string_list(self, value: _str) -> None: ...
    def set_vocab(self, value: dict[_str, _int]) -> None: ...
    @property
    def shape(self) -> paddle.Size: ...
    @property
    def size(self) -> _int: ...
    def sparse_dim(self) -> _int: ...
    @property
    def stop_gradient(self) -> _bool: ...
    @stop_gradient.setter
    def stop_gradient(self, value: _bool) -> None: ...
    @property
    def strides(self) -> list[_int]: ...
    @property
    def type(self) -> Any: ...

    # virtual methods
    def __iter__(self) -> Iterator[Tensor]: ...  # For iterating over the tensor

    # private methods
    def _grad_ivar(self) -> Tensor | None: ...

    # annotation: ${tensor_alias}

class Tensor(AbstractTensor, AbstractEagerParamBase):
    # annotation: ${tensor_docstring}

    __qualname__: Literal["Tensor"]

# annotation: ${tensor_end}
